import requests
import time
import json
import sys
import os
from loguru import logger
sys.path.append(os.path.abspath('../LLM_APIs'))
from response_process import eval_response, process_llm_history


class DeepSeekClient:
    def __init__(self, api_key, model):
        self.ds_endpoint = "https://api.deepseek.com/v1/chat/completions"
        self.r1_endpoint = 'xxx'
        self.model = model
        self.timeout_setting = {
            "connect": 600,   # 连接超时时间(秒)
            "read": 600,      # 读取超时时间(秒)
            "retries": 3,     # 最大重试次数
            "backoff_factor": 0.5  # 指数退避因子
        }
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        })

    def _handle_timeout_retry(self, func, *args, **kwargs): # 带指数退避的重试机制
        for attempt in range(self.timeout_setting["retries"]):
            try:
                return func(*args, **kwargs)
            except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e:
                if attempt == self.timeout_setting["retries"] - 1:
                    raise
                sleep_time = self.timeout_setting["backoff_factor"] * (2 ** attempt)
                print(f"Attempt {attempt+1} failed, retrying in {sleep_time}s...")
                time.sleep(sleep_time)
    
    def generate_content(self, prompt, temperature, top_p, max_tokens):
        input_prompt = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": temperature,
            "top_p": top_p,
            "max_tokens": max_tokens
        }
        
        try:
            response = self._handle_timeout_retry(
                self.session.post,
                self.ds_endpoint,
                data=json.dumps(input_prompt),
                timeout=(self.timeout_setting["connect"], self.timeout_setting["read"])
            )
            response.raise_for_status()
            return response.json()["choices"][0]["message"]["content"]
        except Exception as e:
            print(f"API Error: {str(e)}")
            return None
        

def build_client2prompt(model='deepseek-chat', **kwargs):
    try:
        model = kwargs['paras']['model']
    except:
        pass
    
    try:
        api_key = kwargs['config']['DS_V3_APIKEY']
        task = kwargs['paras']['task']
        query = kwargs['paras']['query']
        texts = kwargs['paras']['texts']
        ds_client = DeepSeekClient(api_key, model)
    except Exception as e:
        print('Bad paras inputs... because: {}'.format(e))
        return None
        
    his_record = process_llm_history(kwargs['config'], kwargs['paras'])
    print(his_record)
    temperature = 0.1
    top_p = 0.1
    max_tokens = 2000
    
    if task=='gen-titles-dfs':
        topic = kwargs['paras']['topic']
        doc_type = kwargs['paras']['doc_type']
        down_level = kwargs['paras']['down_level']
        num_subs = kwargs['paras']['num_subs']
        max_tokens = 150

        prompt = f"""你是{doc_type}文件撰写方面的专家，需要你以{topic}为当前章节标题，生成下一层的标题，具体要求如下：
                - 仅生成下{down_level}层标题，不要输出当前章节标题，标题数量不超过{num_subs}个，不要输出额外解释。
                - 输出采用Markdown格式，使用"-"和空格区分标题。
                - 生成标题的主题与意思应该与以下已经生成的提纲有区别和关联，不要生成意思重复的标题:
            """
    
    elif task=='gen-titles-oneoff':
        topic = kwargs['paras']['topic']
        doc_type = kwargs['paras']['doc_type']
        num_subs = kwargs['paras']['num_subs']
        max_depth = kwargs['paras']['max_depth']
        max_tokens = kwargs['paras']['max_tokens']
        num_titles = kwargs['paras']['num_titles']
                                
        prompt = f"""你是{doc_type}文件撰写方面的专家，需要你以{topic}为文件主题，生成完整的文件提纲，具体要求如下：
                - 生成的提纲要具备清晰完整的父-子树形层级结构，层级数为{max_depth} ，每层包含的标题数量不超过{num_subs}个。
                - 生成的提纲层级遵循Markdown格式，即所有标题都使用"-"前缀，并使用"-"和空格区分标题层级。
                - 全部标题数量不少于{num_titles}个，不要生成意思重复的标题，除生成提纲之外不要返回任何其他解释。
                - 提纲内容需参考和包括以下内容\n:
                {texts}
            """

    elif task=='gen-root-titles':
        content = kwargs['paras']['content']
        style = kwargs['paras']['style']
        title = kwargs['paras']['title']
        doc_type = kwargs['paras']['doc_type']
        num_titles = kwargs['paras']['num_titles']
        max_tokens = 500
                                
        prompt = f"""你是{doc_type}文件撰写方面的专家，需要你生成文件提纲，文件标题是{title}。具体要求如下：
                - 生成的投标文件提纲要参考以下招标要求\n：{content}\n
                - 仅生成一级提纲（即最高层级的章节标题），标题数量不超过{num_titles}个。
                - 生成的提纲层级遵循Markdown格式，即所有标题都使用"-"前缀分隔。
                - 生成的风格符合{style}，除生成提纲之外不要返回任何其他解释。
                - 生成提纲可以参考以下样例：
                {texts}
            """

    elif task=='extract-topics':
        doc_type = kwargs['paras']['doc_type']
        num_topics = kwargs['paras']['num_topics']

        prompt = f"""你是{doc_type}文件撰写方面的专家，需要分析以下内容，提取其中涉及的主题，具体要求如下：
                '''
                {texts}
                '''
                - 你自主判断提取的主题数量，但不能超过{num_topics}个。
                - 保证主题涉及的内容尽量区分，不要重叠。
                - 不要输出额外解释。
                - 输出采用Markdown格式，使用"-"和空格区分主题。
            """
        
    elif task=='rewrite-paras': # when generating response using KB
        max_tokens = kwargs['paras']['words']
        topic = kwargs['paras']['topic']
        type = kwargs['paras']['type']
        style = kwargs['paras']['style']
        temperature = 0.3
        
        prompt =  f"""
                你是{type}方面的撰写专家，你的任务是撰写一段专业内容，具体要求如下：
                - 内容要严格局限在{topic}主题之下，主要生成主题无关内容。
                - 内容字数保持在{int(max_tokens)}字左右。
                - 可以使用1) a) 等分小点论述，但是不要使用MARKDOWN符号。
                - 可以参考以下知识，确保内容的具体性和专业性：
                {texts}
            """
        
    elif task=='rewrite-sentence':
        max_tokens = int(1.3*len(texts))
        add_req = kwargs['paras']['add_req']
        prompt =  f"""
                你的任务是重写以下句子，使其符合用户要求{add_req}：
                '''{texts}'''
                你仅返回重写后的句子。不要返回额外内容。
            """

    elif task=='filling-tb': # when generating response using KB
        temperature = 0
        top_p = 0.01
        
        prompt = f"""
                你是一名表格填充人员，你将接收基础资料\n：'''{texts}'''
                请根据已有资料，找到对应查询关键词'''{query}'''的答案。注意以下事项：
                - 如果基础资料中没有包括答案，你的答案是"null"
                - 必须返回JSON字典格式，键为"query"，值为你找到的答案，答案要直接简单，不要添加和解释。
                - 不要输出任何其他解释性信息。
            """
        
    return ds_client, prompt, temperature, top_p, max_tokens


def call_ds(model='deepseek-chat', **kwargs):
    def get_answer(prompt, ds_client, temperature, top_p, max_tokens):
        try:
            logger.debug('prompt_len:{}', len(prompt))
            content = ds_client.generate_content(prompt, temperature, top_p, max_tokens)
            if not content:
                return "内容生成失败，请检查网络连接或API配置"
            answer = eval_response(content)
            return answer
        except Exception as e:
            answer = None
            logger.debug('not answer, response:{}', e)
        return answer

    ds_client, prompt, temperature, top_p, max_tokens = build_client2prompt(model, **kwargs)
    answer = get_answer(prompt, ds_client, temperature, top_p, max_tokens)
    print(answer)
    return answer


if __name__ == "__main__":

    print()